import pymc as pm
import matplotlib.pyplot as plt
import arviz as az

from modelchecking.mrm import simulate_mdp

# 使用之前的蒙特卡洛模拟获取的累计奖励数据
# 假设我们已经从之前的模拟中获得了1000次累计奖励
cumulative_rewards = simulate_mdp(num_simulations=100, max_steps=50)

# 使用 PyMC 构建贝叶斯模型
with pm.Model() as model:
    # 先验分布：假设累计奖励的均值服从正态分布，标准差服从半正态分布
    mu_prior = pm.Normal('mu', mu=0, sigma=10)  # 均值的先验分布
    sigma_prior = pm.HalfNormal('sigma', sigma=10)  # 标准差的先验分布

    # 似然函数：观察到的累计奖励数据假设是正态分布的
    likelihood = pm.Normal('obs', mu=mu_prior, sigma=sigma_prior, observed=cumulative_rewards)

    # 进行贝叶斯推断，采样后验分布
    trace = pm.sample(200, return_inferencedata=True)

# 可视化后验分布
az.plot_trace(trace)
plt.show()

# 查看后验均值和标准差的估计结果
posterior_summary = az.summary(trace, hdi_prob=0.95)
print(posterior_summary)

# 获取后验均值和95%可信区间
mu_posterior_mean = posterior_summary.loc['mu', 'mean']
credible_interval_mu = posterior_summary.loc['mu', ['hdi_2.5%', 'hdi_97.5%']]

sigma_posterior_mean = posterior_summary.loc['sigma', 'mean']
credible_interval_sigma = posterior_summary.loc['sigma', ['hdi_2.5%', 'hdi_97.5%']]

print(f"后验均值 (mu): {mu_posterior_mean}")
print(f"95% 可信区间 (mu): {credible_interval_mu}")
print(f"后验标准差 (sigma): {sigma_posterior_mean}")
print(f"95% 可信区间 (sigma): {credible_interval_sigma}")
